# ''' Python 代码开始 (已集成可配置坐标系)

import cv2
import numpy as np
import glob
import yaml
import os
import argparse
import abc

# ... [此处省略未改变的 CalibrationPattern 类的代码] ...
class CalibrationPattern(abc.ABC):
    def __init__(self, params): self.params = params; self._validate_params()
    @abc.abstractmethod
    def _validate_params(self): pass
    @abc.abstractmethod
    def get_object_points(self): pass
    @abc.abstractmethod
    def find_corners(self, image): pass
    def draw_corners(self, image, corners, ret): return cv2.drawChessboardCorners(image, self.pattern_size, corners, ret)
class CheckerboardPattern(CalibrationPattern):
    def _validate_params(self):
        if not all(k in self.params for k in ['cols', 'rows', 'square_size']): raise ValueError("棋盘格模式需要 'cols', 'rows', 'square_size' 参数.")
        self.pattern_size = (self.params['cols'], self.params['rows']); self.square_size = self.params['square_size']
    def get_object_points(self):
        objp = np.zeros((self.pattern_size[0] * self.pattern_size[1], 3), np.float32)
        objp[:, :2] = np.mgrid[0:self.pattern_size[0], 0:self.pattern_size[1]].T.reshape(-1, 2)
        return objp * self.square_size
    def find_corners(self, image):
        flags = cv2.CALIB_CB_ADAPTIVE_THRESH + cv2.CALIB_CB_NORMALIZE_IMAGE
        ret, corners = cv2.findChessboardCorners(image, self.pattern_size, flags=flags)
        if not ret: ret, corners = cv2.findChessboardCorners(image, self.pattern_size, cv2.CALIB_CB_FAST_CHECK)
        return ret, corners
class CirclesGridPattern(CalibrationPattern):
    def _validate_params(self):
        if not all(k in self.params for k in ['cols', 'rows', 'spacing']): raise ValueError("对称圆形网格模式需要 'cols', 'rows', 'spacing' 参数.")
        self.pattern_size = (self.params['cols'], self.params['rows']); self.spacing = self.params['spacing']
    def get_object_points(self):
        objp = np.zeros((self.pattern_size[0] * self.pattern_size[1], 3), np.float32)
        objp[:, :2] = np.mgrid[0:self.pattern_size[0], 0:self.pattern_size[1]].T.reshape(-1, 2)
        return objp * self.spacing
    def find_corners(self, image):
        return cv2.findCirclesGrid(image, self.pattern_size, flags=cv2.CALIB_CB_SYMMETRIC_GRID)
class AsymmetricCirclesGridPattern(CalibrationPattern):
    def _validate_params(self):
        if not all(k in self.params for k in ['cols', 'rows', 'spacing']): raise ValueError("非对称圆形网格模式需要 'cols', 'rows', 'spacing' 参数.")
        self.pattern_size = (self.params['cols'], self.params['rows']); self.spacing = self.params['spacing']
    def get_object_points(self):
        objp = np.zeros((self.pattern_size[0] * self.pattern_size[1], 3), np.float32)
        for i in range(self.pattern_size[1]):
            for j in range(self.pattern_size[0]): objp[i * self.pattern_size[0] + j] = ((2*j + i%2) * self.spacing, i * self.spacing, 0)
        return objp
    def find_corners(self, image):
        flags = cv2.CALIB_CB_ASYMMETRIC_GRID + cv2.CALIB_CB_CLUSTERING
        return cv2.findCirclesGrid(image, self.pattern_size, flags=flags)


class CameraCalibrator:
    # ==================== 关键修改点 1: 定义坐标系预设 ====================
    # 定义从 OpenCV标准坐标系 (X-右, Y-下, Z-前) 到目标坐标系的转换矩阵
    # R_target = T @ R_opencv
    COORDINATE_SYSTEM_TRANSFORMS = {
        # 默认: OpenCV标准, 转换矩阵为单位矩阵
        'opencv_standard': np.identity(3),
        # OpenGL标准 (X-右, Y-上, Z-后)
        'opengl_standard': np.array([
            [1, 0, 0],
            [0, -1, 0],
            [0, 0, -1]
        ]),
        # 机器人标准 (X-前, Y-左, Z-上), e.g., ROS REP 103
        'robot_x_front_y_left_z_up': np.array([
            [0, 0, 1],   # 新X轴 = 旧Z轴
            [-1, 0, 0],  # 新Y轴 = 旧-X轴
            [0, -1, 0]   # 新Z轴 = 旧-Y轴
        ]),
        # Unreal Engine 标准 (X-前, Y-右, Z-上), 左手系
        'unreal_engine': np.array([
            [0, 0, 1],   # 新X轴 = 旧Z轴
            [1, 0, 0],   # 新Y轴 = 旧X轴
            [0, -1, 0]   # 新Z轴 = 旧-Y轴
        ])
    }
    # ====================================================================

    PATTERN_HANDLERS = {
        'checkerboard': CheckerboardPattern,
        'circles': CirclesGridPattern,
        'asymmetric_circles': AsymmetricCirclesGridPattern
    }

    def __init__(self, config_path):
        print("INFO: 初始化相机标定器...")
        with open(config_path, 'r', encoding='utf-8') as f:
            self.config = yaml.safe_load(f)
        
        self.output_path = self.config['output_path']
        self.pattern_type = self.config['pattern_type']
        self.pattern_params = self.config['pattern_params']
        
        handler_class = self.PATTERN_HANDLERS.get(self.pattern_type)
        if not handler_class: raise ValueError(f"错误: 不支持的标定板类型 '{self.pattern_type}'.")
        self.pattern_handler = handler_class(self.pattern_params)
        print(f"INFO: 已选择标定板类型: '{self.pattern_type}'")

        if not os.path.exists(self.output_path):
            os.makedirs(self.output_path)
            print(f"INFO: 已创建输出目录: {self.output_path}")

        # ==================== 关键修改点 2: 加载坐标系配置 ====================
        coord_system_name = self.config.get('camera_coordinate_system', 'opencv_standard')
        self.coord_transform_matrix = self.COORDINATE_SYSTEM_TRANSFORMS.get(coord_system_name)
        if self.coord_transform_matrix is None:
            raise ValueError(f"错误: 不支持的坐标系名称 '{coord_system_name}'. "
                             f"可用预设: {list(self.COORDINATE_SYSTEM_TRANSFORMS.keys())}")
        self.coord_system_name = coord_system_name
        print(f"INFO: 将使用坐标系 '{self.coord_system_name}' 输出位姿结果。")
        # ======================================================================

    # calibrate 和 _evaluate_and_save_results 方法无变化, 已折叠
    def calibrate(self):
        print("\nINFO: ============== 开始标定模式 ==============")
        self.image_path = self.config['image_path']
        self.save_verification = self.config['save_verification_images']
        self.save_format = self.config['save_format']
        if self.save_verification:
            self.verification_path = os.path.join(self.output_path, 'verification_outputs')
            if not os.path.exists(self.verification_path): os.makedirs(self.verification_path)
        objp = self.pattern_handler.get_object_points()
        objpoints, imgpoints, valid_image_files = [], [], []
        images = glob.glob(self.image_path)
        if not images: print(f"ERROR: 在路径 '{self.image_path}' 未找到任何图像."); return
        print(f"INFO: 找到 {len(images)} 张图像进行处理.")
        gray = None
        for i, fname in enumerate(sorted(images)):
            img = cv2.imread(fname)
            if img is None: continue
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            print(f"INFO: 正在处理图像 {i+1}/{len(images)}: {os.path.basename(fname)}")
            ret, corners = self.pattern_handler.find_corners(gray)
            if ret:
                objpoints.append(objp)
                imgpoints.append(corners)
                valid_image_files.append(fname)
                print(f"  - 成功找到角点.")
                if self.save_verification:
                    drawn_img = self.pattern_handler.draw_corners(img.copy(), corners, ret)
                    cv2.imwrite(os.path.join(self.verification_path, f'corners_found_{os.path.basename(fname)}'), drawn_img)
            else:
                print(f"  - 未能找到角点, 跳过此图像.")
        if not objpoints: print("ERROR: 所有图像均未能成功检测到角点，标定失败。"); return
        print(f"\nINFO: 使用 {len(objpoints)} 张有效图像进行标定...")
        if gray is None: print("ERROR: 未能成功加载任何图像的灰度图。"); return
        h, w = gray.shape[:2]
        calib_flags = 0
        flag_str = self.config.get('calibration_flags')
        if flag_str and hasattr(cv2, 'CALIB_' + flag_str): calib_flags = getattr(cv2, 'CALIB_' + flag_str)
        ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, (w, h), None, None, flags=calib_flags)
        if not ret: print("ERROR: cv2.calibrateCamera 函数返回失败."); return
        print("\n----------- 标定结果 -----------")
        print(f"相机内参矩阵 (Camera Matrix):\n{mtx}\n\n畸变系数 (Distortion Coefficients):\n{dist}")
        print("---------------------------------\n")
        self._evaluate_and_save_results(mtx, dist, rvecs, tvecs, objpoints, imgpoints, valid_image_files)
    def _evaluate_and_save_results(self, mtx, dist, rvecs, tvecs, objpoints, imgpoints, valid_image_files):
        mean_error = sum(cv2.norm(imgpoints[i], cv2.projectPoints(objpoints[i], rvecs[i], tvecs[i], mtx, dist)[0], cv2.NORM_L2) / len(imgpoints[i]) for i in range(len(objpoints)))
        reprojection_error = mean_error / len(objpoints)
        print(f"INFO: 平均重投影误差 (Mean Reprojection Error): {reprojection_error:.4f} 像素")
        with open(os.path.join(self.output_path, 'reprojection_error.txt'), 'w') as f: f.write(f"Mean Reprojection Error: {reprojection_error}\n")
        if self.save_format == 'npz':
            save_path = os.path.join(self.output_path, 'calibration_results.npz')
            np.savez(save_path, mtx=mtx, dist=dist, rvecs=np.array(rvecs, dtype=object), tvecs=np.array(tvecs, dtype=object))
        elif self.save_format == 'yaml':
            save_path = os.path.join(self.output_path, 'calibration_results.yaml')
            data = {'camera_matrix': mtx.tolist(), 'distortion_coefficients': dist.tolist()}
            with open(save_path, 'w', encoding='utf-8') as f: yaml.dump(data, f, default_flow_style=False)
        print(f"INFO: 标定结果已保存至: {save_path}")
        if not self.save_verification: print("\n标定流程完成！"); return
        print("INFO: 正在生成并保存验证图像...")
        alpha = self.config.get('undistort_alpha', 0.0)
        for i, fname in enumerate(valid_image_files):
            img = cv2.imread(fname)
            h, w = img.shape[:2]
            newcameramtx, roi = cv2.getOptimalNewCameraMatrix(mtx, dist, (w, h), alpha, (w, h))
            dst = cv2.undistort(img, mtx, dist, None, newcameramtx)
            axis_img = img.copy()
            unit_size = self.pattern_handler.params.get('square_size', self.pattern_handler.params.get('spacing', 30))
            cv2.drawFrameAxes(axis_img, mtx, dist, rvecs[i], tvecs[i], 3 * unit_size)
            cv2.imwrite(os.path.join(self.verification_path, f'axes_on_board_{os.path.basename(fname)}'), axis_img)
        print("INFO: 所有验证图像已保存完毕.\n\n标定流程完成！")


    def load_calibration_results(self, calib_file_path):
        if not os.path.exists(calib_file_path): raise FileNotFoundError(f"标定文件未找到: {calib_file_path}")
        _, ext = os.path.splitext(calib_file_path)
        if ext == '.npz':
            with np.load(calib_file_path) as data: mtx, dist = data['mtx'], data['dist']
        elif ext in ['.yaml', '.yml']:
            with open(calib_file_path, 'r', encoding='utf-8') as f:
                data = yaml.safe_load(f)
                mtx, dist = np.array(data['camera_matrix']), np.array(data['distortion_coefficients'])
        else: raise ValueError(f"不支持的标定文件格式: {ext}。")
        print(f"INFO: 成功从 {calib_file_path} 加载标定数据。")
        return mtx, dist

    def evaluate_pose(self, image_path, calib_file_path):
        print("\nINFO: ============== 开始位姿估计模式 ==============")
        if not os.path.exists(image_path): print(f"ERROR: 图像 '{image_path}' 不存在."); return
        try:
            mtx, dist = self.load_calibration_results(calib_file_path)
        except (FileNotFoundError, ValueError, KeyError) as e: print(f"ERROR: 加载标定文件失败. {e}"); return

        img = cv2.imread(image_path)
        if img is None: print(f"ERROR: 无法读取图像: {image_path}"); return
        
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        objp = self.pattern_handler.get_object_points()

        print(f"INFO: 正在图像 '{os.path.basename(image_path)}' 中寻找角点...")
        ret, corners = self.pattern_handler.find_corners(gray)

        if not ret: print("ERROR: 未能在此图像中找到角点, 无法估计位姿。"); return
        print("INFO: 成功找到角点.")
        
        ret, rvec_opencv, tvec_opencv = cv2.solvePnP(objp, corners, mtx, dist)
        if not ret: print("ERROR: cv2.solvePnP 函数返回失败."); return

        # ==================== 关键修改点 3: 应用坐标系转换 ====================
        
        # 1. 将OpenCV的旋转向量转为旋转矩阵
        R_world_to_opencv, _ = cv2.Rodrigues(rvec_opencv)

        # 2. 获取从OpenCV坐标系到用户定义坐标系的转换矩阵
        T_opencv_to_user = self.coord_transform_matrix

        # 3. 计算用户坐标系下的旋转和平移
        #    R_world_to_user = T_opencv_to_user * R_world_to_opencv
        R_world_to_user = T_opencv_to_user @ R_world_to_opencv
        tvec_user = T_opencv_to_user @ tvec_opencv

        # 4. 将用户坐标系下的旋转矩阵转回旋转向量，以便打印
        rvec_user, _ = cv2.Rodrigues(R_world_to_user)
        
        # 5. 标定板坐标轴在用户相机坐标系下的方向矢量 (即 R_world_to_user 的列向量)
        x_axis_in_user_cam = R_world_to_user[:, 0]
        y_axis_in_user_cam = R_world_to_user[:, 1]
        z_axis_in_user_cam = R_world_to_user[:, 2]

        print(f"\n----------- 位姿估计结果 (坐标系: {self.coord_system_name}) -----------")
        print(f"旋转向量 (rvec):\n{rvec_user.flatten()}")
        print(f"\n平移向量 (tvec) (单位: mm):\n{tvec_user.flatten()}")
        
        print("\n--- 标定板坐标轴在相机坐标系下的方向矢量 ---")
        print(f"X 轴 (标定板) 方向: {np.round(x_axis_in_user_cam, 4)}")
        print(f"Y 轴 (标定板) 方向: {np.round(y_axis_in_user_cam, 4)}")
        print(f"Z 轴 (标定板) 方向: {np.round(z_axis_in_user_cam, 4)}")
        print("-------------------------------------------------------------------\n")

        # ========================================================================

        print("INFO: 正在生成可视化结果...")
        unit_size = self.pattern_handler.params.get('square_size', self.pattern_handler.params.get('spacing', 30))
        # 可视化仍然使用OpenCV的drawFrameAxes, 它只认识OpenCV坐标系, 所以传入原始的rvec/tvec
        axis_img = cv2.drawFrameAxes(img.copy(), mtx, dist, rvec_opencv, tvec_opencv, 3 * unit_size)

        output_filename = f'pose_estimation_on_{os.path.basename(image_path)}'
        save_path = os.path.join(self.output_path, output_filename)
        cv2.imwrite(save_path, axis_img)
        print(f"INFO: 位姿可视化图像已保存至: {save_path} (注意: 图像上的坐标轴总是以OpenCV标准绘制)")

        print("\n位姿估计完成！")

# ===================================================================
# 主函数入口 (无变化)
# ===================================================================
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="相机标定与位姿估计脚本.")
    parser.add_argument('--config', type=str, default='camera_config.yaml', help='配置文件的路径')
    parser.add_argument('--mode', type=str, default='calibrate', choices=['calibrate', 'pose'], help="运行模式")
    parser.add_argument('--image', type=str, help="[位姿估计模式] 用于估计位姿的单张图像路径.")
    parser.add_argument('--calib', type=str, default='output/calibration_results.npz', help="[位姿估计模式] 已保存的标定结果文件路径")
    args = parser.parse_args()
    if not os.path.exists(args.config):
        print(f"FATAL: 配置文件 '{args.config}' 不存在."); exit()
    try:
        calibrator = CameraCalibrator(args.config)
        if args.mode == 'calibrate':
            calibrator.calibrate()
        elif args.mode == 'pose':
            if not args.image: parser.error("--mode 'pose' 需要提供 --image 参数。")
            calibrator.evaluate_pose(args.image, args.calib)
    except (ValueError, FileNotFoundError, KeyError, TypeError) as e:
        print(f"\nERROR: 处理失败. 详情: {e}")

# ''' Python代码结束